import requests  # 用于发送HTTP请求
import json  # 用于解析JSON数据


def chatAiPost(url, key, data, OnDataProcess, OnDataFinish):
    """
    发送流式请求到AI聊天接口，并处理返回的数据流

    :param url: 请求的URL，字符串类型
    :param key: 鉴权密钥，字符串类型
    :param data: 用户传入的消息数据，将被封装到请求体中
    :param OnDataProcess: 数据接收过程的回调函数，每次接收到内容时调用
    :param OnDataFinish: 数据接收完成的回调函数，所有内容接收完成后调用
    """
    # 构造符合要求的请求体
    payload = {
        "stream": False,
        "detail": False,
        "messages": data  # 用户消息直接作为messages字段的值
    }

    # 设置请求头，包含鉴权信息
    headers = {
        "Authorization": f"Bearer {key}",  # 添加Bearer Token
        "Content-Type": "application/json"
    }

    try:
        # 发送POST请求，启用流式传输模式
        response = requests.post(
            url,
            json=payload,
            headers=headers,
            stream=True  # 重要：启用流式传输
        )
        response.raise_for_status()  # 如果响应状态码异常则抛出错误
    except requests.exceptions.RequestException as e:
        print(f"请求发生错误: {e}")
        return  # 可在此添加更详细的错误处理逻辑

    collected_content = []  # 用于收集所有有效内容片段

    # 逐行处理流式响应内容
    for line in response.iter_lines():
        if line:  # 过滤掉保持连接的空行
            decoded_line = line.decode('utf-8').strip()

            # 仅处理以"data: "开头的有效数据行
            if decoded_line.startswith('data: '):
                event_data = decoded_line[len('data: '):]  # 提取实际数据部分

                # 判断是否结束标记
                if event_data == '[DONE]':
                    break  # 终止数据接收循环

                try:
                    data_json = json.loads(event_data)

                    # 提取第一个choice的delta.content
                    if data_json.get('choices'):
                        first_choice = data_json['choices'][0]
                        content = first_choice.get('delta', {}).get('content')

                        # 过滤掉content为null的情况（如结束标记）
                        if content is not None:
                            collected_content.append(content)
                            OnDataProcess(content)  # 实时回调单个内容片段
                except json.JSONDecodeError:
                    pass  # 忽略非JSON格式的异常数据

    # 所有数据接收完成后，拼接完整内容并触发完成回调
    full_response = ''.join(collected_content)
    OnDataFinish(full_response)


# 以下是使用示例（需根据实际情况实现回调函数）
if __name__ == "__main__":
    def example_process(data):
        print(f"收到片段: {data}")


    def example_finish(data):
        print(f"完整响应: {data}")


    chatAiPost(
        url="https://cloud.fastgpt.cn/api/v1/chat/completions",
        key="fastgpt-vQUr1PVP035M22uYV8hUhHYljFXV7FF9o4fDCaHrP2gnnjUB676WAOs9Gn",
        data=[{"role": "user", "content": "你好"}],
        OnDataProcess=example_process,
        OnDataFinish=example_finish
    )